package com.wzw.expression.record.support;

import com.alibaba.ttl.TransmittableThreadLocal;
import lombok.NoArgsConstructor;

import java.io.Serializable;
import java.util.Objects;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author Wangzhiwen
 */
public class IRecordTraces {

    private static final ThreadLocal<ThreadRecord> RECORD = new TransmittableThreadLocal<>() {
        @Override
        protected ThreadRecord initialValue() {
            return new ThreadRecord(UUID.randomUUID().toString(), new AtomicInteger(0));
        }
    };

    public static String getTraceId() {
        return RECORD.get().traceId;
    }

    public static boolean done() {
        return RECORD.get().count.get() == 0;
    }

    public static void increment() {
        RECORD.get().count.incrementAndGet();
    }

    public static void decrement() {
        RECORD.get().count.decrementAndGet();
    }

    public static void remove() {
        RECORD.remove();
    }

    @NoArgsConstructor
    private static class ThreadRecord implements Serializable {

        private String traceId;

        private AtomicInteger count;

        public ThreadRecord(String traceId, AtomicInteger count) {
            this.traceId = traceId;
            this.count = count;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) {
                return true;
            }
            if (!(o instanceof ThreadRecord)) {
                return false;
            }
            ThreadRecord that = (ThreadRecord) o;
            return Objects.equals(traceId, that.traceId);
        }

        @Override
        public int hashCode() {
            return Objects.hash(traceId);
        }
    }

}
